"""SynergyNet loss definition."""
import math
import mindspore
import mindspore.ops as ops
from mindspore.nn.loss.loss import LossBase
import mindspore as ms
import mindspore.nn as nn
from mindspore.common.tensor import Tensor
from mindspore import context
from mindspore.context import ParallelMode
from mindspore.parallel._auto_parallel_context import auto_parallel_context
from mindspore.communication.management import get_group_size
import numpy as np
import mindspore.numpy as mnp
from mindspore.ops import operations as P
from mindspore.ops import functional as F
from mindspore.ops import composite as C


class WingLoss(LossBase):
    def __init__(self, omega=10.0, epsilon=2):
        super(WingLoss, self).__init__()
        self.omega = omega
        self.epsilon = epsilon
        self.long_term = math.log(1 + self.omega / self.epsilon)

    def construct(self, pred, target, kp=False):
        n_points = pred.shape[2]
        pred = pred.transpose(0, 2, 1)
        pred = pred.view(-1, 3 * n_points)
        target = target.transpose(0, 2, 1)
        target = target.view(-1, 3 * n_points)
        y = target
        y_hat = pred

        delta_y = (y - y_hat).abs()
        delta_y1 = ops.masked_select(delta_y, delta_y < self.omega)
        delta_y2 = ops.masked_select(delta_y, delta_y >= self.omega)
        log = ops.Log()
        loss1 = self.omega * log(1 + delta_y1 / self.epsilon)
        C = self.omega - self.omega * self.long_term
        loss2 = delta_y2 - C
        return (loss1.sum() + loss2.sum()) / (delta_y.shape[0] * delta_y.shape[1])


class ParamLoss(LossBase):
    """Input and target are all 62-d param"""

    def __init__(self):
        super(ParamLoss, self).__init__()
        self.criterion = nn.MSELoss(reduction='none')  # 均方误差
        self.sqrt = ops.Sqrt()

    def construct(self, input, target, mode='normal'):
        # input = input, dtype=mindspore.float32)
        # target = Tensor(target, dtype=mindspore.float32)
        out = 0
        if mode == 'normal':
            input1 = input[:, :12]
            target1 = target[:, :12]
            loss1 = self.criterion(input1, target1).astype(mindspore.float32)
            loss1 = mindspore.numpy.mean(loss1, axis=1)
            loss2 = self.criterion(input1, target1).astype(mindspore.float32)
            loss2 = mindspore.numpy.mean(loss2, axis=1)
            loss = loss1 + loss2
            out = self.sqrt(loss)
        elif mode == 'only_3dmm':
            loss = self.criterion(input[:, :50], target[:, 12:62]).mean(axis=1)
            out = self.sqrt(loss)
        return out


class Synergynet_Loss(LossBase):
    """SynergynetLosses"""

    def __init__(self, reduction="mean"):
        super(Synergynet_Loss, self).__init__(reduction)
        self.LMKLoss_3D = WingLoss()
        self.ParamLoss = ParamLoss()

    def construct(self, _3D_attr, _3D_attr_GT, vertex_lmk1, vertex_GT_lmk, vertex_lmk2, _3D_attr_S2):
        loss_LMK_f0 = 0.05 * self.LMKLoss_3D(vertex_lmk1, vertex_GT_lmk, kp=True)
        loss_Param_In = 0.02 * self.ParamLoss(_3D_attr, _3D_attr_GT)
        loss_LMK_pointNet = 0.05 * self.LMKLoss_3D(vertex_lmk2, vertex_GT_lmk, kp=True)
        loss_Param_S2 = 0.02 * self.ParamLoss(_3D_attr_S2, _3D_attr_GT, mode='only_3dmm')
        loss_Param_S1S2 = 0.001 * self.ParamLoss(_3D_attr_S2, _3D_attr, mode='only_3dmm')

        total_loss = loss_LMK_f0.mean() + loss_Param_In.mean() + loss_LMK_pointNet.mean() + loss_Param_S2.mean() + loss_Param_S1S2.mean()
        return total_loss


class CustomWithLossCell(nn.Cell):
    """SynergyNet loss"""

    def __init__(self, backbone, loss_fn):
        super(CustomWithLossCell, self).__init__()
        self._backbone = backbone
        self._loss_fn = loss_fn

    def construct(self, x, target):
        _3D_attr, _3D_attr_GT, vertex_lmk1, vertex_GT_lmk, vertex_lmk2, _3D_attr_S2 = self._backbone(x, target)
        return self._loss_fn(_3D_attr, _3D_attr_GT, vertex_lmk1, vertex_GT_lmk, vertex_lmk2, _3D_attr_S2)


class AverageMeter:
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count
